import torch
import torch.nn as nn
import torch.nn.functional as F

class ActionClassificationModel(nn.Module):

    def __init__(self, num_classes=5, input_channels=3, img_height=112, img_width=112, time_steps=5, hidden_size=256):
        super(ActionClassificationModel, self).__init__()
        self.num_classes = num_classes
        self.input_channels = input_channels
        self.img_height = img_height
        self.img_width = img_width
        self.time_steps = time_steps
        self.hidden_size = hidden_size

        # CNN layers
        self.conv1 = nn.Conv3d(input_channels, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm3d(32)
        self.conv2 = nn.Conv3d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm3d(64)
        self.conv3 = nn.Conv3d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm3d(128)

        # LSTM layers
        self.lstm1 = nn.LSTM(input_size=128 * (img_height // 8) * (img_width // 8), hidden_size=hidden_size,
                             batch_first=True)
        self.lstm2 = nn.LSTM(input_size=hidden_size, hidden_size=hidden_size, batch_first=True)

        # Fully-connected layers
        self.fc1 = nn.Linear(hidden_size, 64)
        self.fc2 = nn.Linear(64, num_classes)

    def forward(self, x):
        # x.shape: [batch_size, time_steps, input_channels, img_height, img_width]
        batch_size, time_steps, input_channels, img_height, img_width = x.shape
        x = x.view(batch_size ,input_channels,time_steps,  img_height, img_width)  # [batch_size * time_steps, input_channels, img_height, img_width]

        # CNN
        x = self.bn1(F.relu(self.conv1(x)))
        x = F.max_pool3d(x, (1, 2, 2))
        x = self.bn2(F.relu(self.conv2(x)))
        x = F.max_pool3d(x, (1, 2, 2))
        x = self.bn3(F.relu(self.conv3(x)))
        x = F.max_pool3d(x, (1, 2, 2))
        x = x.view(batch_size * time_steps, 128 * (self.img_height // 8) * (self.img_width // 8))  # [batch_size * time_steps, 128 * (self.img_height // 8) * (self.img_width // 8)]

        # LSTM
        x = x.view(batch_size, time_steps, -1)  # [batch_size, time_steps, 128 * (self.img_height // 8) * (self.img_width // 8)]
        h0 = torch.zeros(1, batch_size, self.hidden_size).to(x.device)
        c0 = torch.zeros(1, batch_size, self.hidden_size).to(x.device)
        out, (hn, cn) = self.lstm1(x, (h0, c0))
        out, (hn, cn) = self.lstm2(out, (hn, cn))

        # Fully-connected
        out = F.relu(self.fc1(out[:, -1, :]))
        out = self.fc2(out)
        return out

